package com.example.filetransfer.server;

import com.example.filetransfer.util.Base64Util;
import java.io.*;
import java.net.ServerSocket;
import java.net.Socket;
import java.util.Base64;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.function.Consumer;

public class FileTransferServer {
    private final int port;
    private ServerSocket serverSocket;
    private ExecutorService executorService;
    private volatile boolean running;
    private Consumer<String> messageHandler;

    public FileTransferServer(int port) {
        this.port = port;
        this.executorService = Executors.newCachedThreadPool();
    }

    public void setMessageHandler(Consumer<String> messageHandler) {
        this.messageHandler = messageHandler;
    }

    public void start() throws IOException {
        serverSocket = new ServerSocket(port);
        running = true;
        log("服务器启动，监听端口: " + port);

        executorService.submit(() -> {
            try {
                while (running) {
                    Socket clientSocket = serverSocket.accept();
                    handleClient(clientSocket);
                }
            } catch (IOException e) {
                if (running) {
                    log("接受客户端连接时发生错误: " + e.getMessage());
                }
            }
        });
    }

    public void stop() {
        running = false;
        try {
            if (serverSocket != null && !serverSocket.isClosed()) {
                serverSocket.close();
            }
        } catch (IOException e) {
            log("关闭服务器时发生错误: " + e.getMessage());
        }
        executorService.shutdown();
    }

    private void handleClient(Socket clientSocket) {
        executorService.submit(() -> {
            try {
                log("客户端已连接: " + clientSocket.getInetAddress().getHostAddress());
                DataInputStream in = new DataInputStream(clientSocket.getInputStream());
                DataOutputStream out = new DataOutputStream(clientSocket.getOutputStream());

                while (running) {
                    int type = in.readInt();
                    switch (type) {
                        case 1: // 文本消息
                            handleTextMessage(in, out);
                            break;
                        case 2: // 文件传输
                            handleFileTransfer(in, out);
                            break;
                        default:
                            log("未知的消息类型: " + type);
                            return;
                    }
                }
            } catch (IOException e) {
                if (running) {
                    log("处理客户端时发生错误: " + e.getMessage());
                }
            } finally {
                try {
                    clientSocket.close();
                } catch (IOException e) {
                    log("关闭客户端连接时发生错误: " + e.getMessage());
                }
            }
        });
    }

    private void handleTextMessage(DataInputStream in, DataOutputStream out) throws IOException {
        int length = in.readInt();
        byte[] messageBytes = new byte[length];
        in.readFully(messageBytes);
        String encodedMessage = new String(messageBytes, "UTF-8");
        // 解码Base64消息
        String message = Base64Util.decodeString(encodedMessage);
        log("收到消息: " + message);

        // 发送确认（也使用Base64编码）
        String encodedResponse = Base64Util.encodeString("消息已接收");
        out.writeInt(1);
        out.writeInt(encodedResponse.getBytes("UTF-8").length);
        out.write(encodedResponse.getBytes("UTF-8"));
    }

    private void handleFileTransfer(DataInputStream in, DataOutputStream out) throws IOException {
        // 读取文件名长度
        int fileNameLength = in.readInt();
        byte[] fileNameBytes = new byte[fileNameLength];
        in.readFully(fileNameBytes);
        String encodedFileName = new String(fileNameBytes, "UTF-8");
        // 解码Base64文件名
        String fileName = Base64Util.decodeString(encodedFileName);

        // 读取文件大小
        long fileSize = in.readLong();
        log("开始接收文件: " + fileName + " (" + fileSize + " bytes)");

        // 创建保存目录
        File uploadDir = new File("uploads");
        if (!uploadDir.exists()) {
            uploadDir.mkdirs();
        }

        // 保存文件
        File file = new File(uploadDir, fileName);
        try (FileOutputStream fos = new FileOutputStream(file)) {
            byte[] buffer = new byte[8192];
            long remaining = fileSize;
            while (remaining > 0) {
                int read = in.read(buffer, 0, (int) Math.min(buffer.length, remaining));
                if (read == -1) {
                    throw new IOException("文件传输中断");
                }
                // 解码Base64文件内容
                byte[] decodedBytes = Base64.getDecoder().decode(
                    java.util.Arrays.copyOfRange(buffer, 0, read)
                );
                fos.write(decodedBytes);
                remaining -= read;
            }
        }

        log("文件接收完成: " + fileName);

        // 发送确认（使用Base64编码）
        String encodedResponse = Base64Util.encodeString("文件已接收");
        out.writeInt(2);
        out.writeInt(encodedResponse.getBytes("UTF-8").length);
        out.write(encodedResponse.getBytes("UTF-8"));
    }

    private void log(String message) {
        if (messageHandler != null) {
            messageHandler.accept(message);
        }
    }
} 